#include "rewine/impl/CDsoImage.hpp"

#include "dso/DsoLoader.hpp"
#include "infra/Shelter.hpp"
#include "rewine/LoggerSingleton.hpp"
#include "rewine/impl/CDsoExport.hpp"

#include <link.h>
#include <elf.h>

#define ELFW(type)	_ElfW (ELF, __ELF_NATIVE_CLASS, type)

using std::shared_ptr;
using std::string;

namespace rewine
{

int
CDsoImage::init(const std::string &exportName, const std::string &fullname)
{
    setFullName(fullname);
    setExportName(exportName);
    return 0;
}

struct dl_iterate_phdr_context_t {
    int found;
    ElfW(Addr) find_start;
    ElfW(Addr) find_end;
    ElfW(Addr) ret_vstart;
    ElfW(Addr) ret_vend;
    ElfW(Addr) ret_fstart;
    ElfW(Addr) ret_fend;
};

static int
dl_find_in_objects(struct dl_phdr_info *info, size_t size, void *ctx0)
{
    auto ctx = (struct dl_iterate_phdr_context_t *)ctx0;

    for (int i = 0; i < info->dlpi_phnum; ++i) {
        auto phdr = info->dlpi_phdr + i;
        ElfW(Addr) seg_start = info->dlpi_addr + phdr->p_vaddr;
        ElfW(Addr) seg_end = seg_start + phdr->p_memsz;
        if (seg_start <= ctx->find_start && ctx->find_end <= seg_end) {
            // segment found
            ctx->ret_vstart = seg_start;
            ctx->ret_vend = seg_end;
            ctx->ret_fstart = phdr->p_offset;
            ctx->ret_fend = ctx->ret_fstart + phdr->p_filesz;
            ctx->found = 1;
            return 1;
        }
    }
    return 0;
}

int
CDsoImage::read_dynamic_symbol_count()
{
    struct link_map *l = (struct link_map *) rewine_dlinfo_map(*m_Handle);
    if (!l) {
        auto sErr = rewine_dlerror();
        g_Logger->error(__FLF__, "{} dlinfo err: {}", exportName().c_str(), sErr);
        return -1;
    }

    const ElfW(Sym) *symtab = nullptr;
    const char *strtab = nullptr;
    for (ElfW(Dyn) *ld = l->l_ld; ld->d_tag != DT_NULL; ++ld) {
        switch (ld->d_tag) {
            case DT_SYMTAB: {
                symtab = (ElfW(Sym) *) (ld->d_un.d_ptr);
                break;
            }
            case DT_STRTAB: {
                strtab = (const char *) (ld->d_un.d_ptr);
                break;
            }
        }
    }
    if (!(symtab && strtab)) {
        g_Logger->info(__FLF__, "{} symtab={} strtab={}", exportName().c_str(), (void *)symtab, (void *)strtab);
        return 0;
    }

    //TODO: a dangerous assumption: the .dymsym is followed by the .dymstr and no padding between them
    size_t symcnt = ((char *)strtab - (char *)symtab) / sizeof(ElfW(Sym));
    g_Logger->debug(__FLF__, "has {} symbols", symcnt);

    m_DynamicSymbolCount = symcnt;
    m_base = (void *)l->l_addr;
    m_symtab = symtab;
    m_strtab = strtab;
    return 0;
}

int
CDsoImage::load()
{
    setIsLoading();
    Shelter shlLoadState;
    shlLoadState.add([this]() {
        setIsLoaded(false);
    });

    m_Handle = rewine_dlopen(fullname().c_str(), true, true);
    if (!m_Handle) {
        auto sErr = rewine_dlerror();
        g_Logger->error(__FLF__, "{} dlopen {} err: {}", exportName().c_str(), fullname().c_str(), sErr);
        return -1;
    }

    int ret;
    ret = read_dynamic_symbol_count();
    if (ret) {
        return -2;
    }

    shlLoadState.toggle(false);
    setIsLoaded(true);
    return 0;
}

int
CDsoImage::unload()
{
    m_Handle.reset();
    setIsUnloaded();
    return 0;
}

int
CDsoImage::getExportCount() const
{
    // remove the first NULL
    return m_DynamicSymbolCount - 1;
}

CDsoExport *
CDsoImage::getCached(const std::string &name) const
{
    auto it = std::find_if(m_Cache.cbegin(), m_Cache.cend(), [&name](const CDsoImage::map_t::value_type& pair) {
        return name == pair.first;
    });
    if (it != m_Cache.cend()) {
        return it->second.get();
    }
    return nullptr;
}

CDsoExport *
CDsoImage::setCached(const std::string &name, void *value) const
{
    std::shared_ptr<CDsoExport> pExport = CDsoExport::create();
    pExport->setHandle(m_Handle, value);
    m_Cache.emplace(std::make_pair(name, pExport));
    return pExport.get();
}

IExport *
CDsoImage::getExportByIndex(unsigned int index) const
{
    auto psym = ((ElfW(Sym) *) m_symtab) + (1 + index);
    std::string name = std::string(m_strtab + psym->st_name);
    auto pExport = getCached(name);
    if (pExport) {
        return pExport;
    }
    if (ELFW(ST_BIND)(psym->st_info) != STB_GLOBAL) {
        g_Logger->error(__FLF__, "{} symbol {} is not global", exportName().c_str(), name.c_str());
        return nullptr;
    }
    void *value = (void *) ((const char *)m_base + psym->st_value);
    return setCached(name, value);
}

IExport *
CDsoImage::getExportByName(const std::string &name, uint16_t hint) const
{
    auto pExport = getCached(name);
    if (pExport != nullptr) {
        return pExport;
    }
    void *address = rewine_dlsym(*m_Handle, name.c_str());
    if (!address) {
        if (auto sErr = rewine_dlerror()) {
            g_Logger->error(__FLF__, "{} dlsym {} err: {}", exportName().c_str(), name.c_str(), sErr);
            return nullptr;
        }
    }
    return setCached(name, address);
}

} // namespace rewine
